package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/spf13/pflag"
)

const (
	dbIndex        = 4         // 使用的 Redis DB
	key            = "make"    // 操作的 key
	value          = "install" // 写入的 value
	cleanupOnExit  = true      // 退出时是否清理 key
	reportInterval = 1000      // 每1000次操作输出统计
)

var (
	// 统计变量
	writeSuccess uint64
	writeFail    uint64
	readSuccess  uint64
	readFail     uint64

	// 控制协程的上下文
	ctx, cancel = context.WithCancel(context.Background())
	wg          sync.WaitGroup

	// Redis 客户端
	rdb *redis.Client
)

func initRedisClient(redisServer string) *redis.Client {
	addr := fmt.Sprintf("%s:6379", redisServer)
	log.Println(addr)
	return redis.NewClient(&redis.Options{
		Addr:         addr,    // Redis 地址
		Password:     "",      // 密码
		DB:           dbIndex, // DB 编号
		ReadTimeout:  10 * time.Millisecond,
		WriteTimeout: 10 * time.Millisecond,
	})
}

func printProgress() {
	totalWrite := atomic.LoadUint64(&writeSuccess) + atomic.LoadUint64(&writeFail)
	totalRead := atomic.LoadUint64(&readSuccess) + atomic.LoadUint64(&readFail)

	if totalWrite%reportInterval == 0 || totalRead%reportInterval == 0 {
		log.Printf("当前进度 - 写入: %d (成功: %d, 失败: %d), 读取: %d (成功: %d, 失败: %d)\n",
			totalWrite, atomic.LoadUint64(&writeSuccess), atomic.LoadUint64(&writeFail),
			totalRead, atomic.LoadUint64(&readSuccess), atomic.LoadUint64(&readFail))
	}
}

func writer() {
	defer wg.Done()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			err := rdb.Set(ctx, key, value, 0).Err()
			if err != nil {
				atomic.AddUint64(&writeFail, 1)
				log.Printf("写入失败: %v", err)
			} else {
				atomic.AddUint64(&writeSuccess, 1)
			}
			printProgress()
			time.Sleep(10 * time.Millisecond)
		}
	}
}

func reader() {
	defer wg.Done()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			_, err := rdb.Get(ctx, key).Result()
			if err != nil && err != redis.Nil {
				atomic.AddUint64(&readFail, 1)
				log.Printf("读取失败: %v", err)
			} else {
				atomic.AddUint64(&readSuccess, 1)
			}
			printProgress()
			time.Sleep(10 * time.Millisecond)
		}
	}
}

func cleanup() {
	if cleanupOnExit {
		cleanupCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()

		_, err := rdb.Del(cleanupCtx, key).Result()
		if err != nil {
			log.Printf("清理 key 失败: %v", err)
		} else {
			log.Println("成功清理 key")
		}
	}
}

func printStats() {
	totalWrite := writeSuccess + writeFail
	totalRead := readSuccess + readFail

	fmt.Println("\n最终测试结果统计:")
	fmt.Printf("写入操作 - 成功: %d, 失败: %d, 成功率: %.2f%%\n",
		writeSuccess, writeFail, float64(writeSuccess)/float64(totalWrite)*100)
	fmt.Printf("读取操作 - 成功: %d, 失败: %d, 成功率: %.2f%%\n",
		readSuccess, readFail, float64(readSuccess)/float64(totalRead)*100)
}

func main() {
	var redisServer string
	pflag.StringVar(&redisServer, "redis_server", "", "redis server")
	pflag.Parse()
	// 初始化 Redis 客户端
	rdb = initRedisClient(redisServer)
	defer rdb.Close()

	// 测试连接
	_, err := rdb.Ping(ctx).Result()
	if err != nil {
		log.Fatalf("无法连接到 Redis: %v", err)
	}

	// 启动写入协程
	wg.Add(1)
	go writer()

	// 启动读取协程
	wg.Add(1)
	go reader()

	// 设置信号捕获
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// 等待终止信号
	<-sigChan
	fmt.Println("\n接收到终止信号，停止测试...")

	// 执行清理
	cleanup()

	// 取消上下文通知协程停止
	cancel()

	// 等待协程结束
	wg.Wait()

	// 打印统计信息
	printStats()
}
